from torch import nn
import torch
from torchsummary import summary


class BasicConv2d(nn.Module):
    def __init__(self, in_channles, out_channles, **kwargs):
        super().__init__()
        self.conv = nn.Conv2d(in_channles, out_channles, bias=False, **kwargs)
        self.bn = nn.BatchNorm2d(out_channles, eps=0.001)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        return nn.functional.relu(x, inplace=True)


# 得到输入大小不变，通道数为768的特征图。
class InceptionC(nn.Module):
    def __init__(self, in_channels, channels_5x5, conv_block=None):
        super().__init__()
        if conv_block is None:
            conv_block = BasicConv2d
        self.branch1x1 = conv_block(in_channels, 16, kernel_size=1)

        c5 = channels_5x5
        self.branch5x5_1 = conv_block(in_channels, c5, kernel_size=1)
        self.branch5x5_2 = conv_block(c5, c5, kernel_size=(1, 5), padding=(0, 2))
        self.branch5x5_3 = conv_block(c5, 16, kernel_size=(5, 1), padding=(2, 0))

        self.branch5x5dbl_1 = conv_block(in_channels, c5, kernel_size=1)
        self.branch5x5dbl_2 = conv_block(c5, c5, kernel_size=(1, 5), padding=(0, 2))
        self.branch5x5dbl_3 = conv_block(c5, c5, kernel_size=(5, 1), padding=(2, 0))
        self.branch5x5dbl_4 = conv_block(c5, c5, kernel_size=(1, 5), padding=(0, 2))
        self.branch5x5dbl_5 = conv_block(c5, 16, kernel_size=(5, 1), padding=(2, 0))

        self.branch_pool = conv_block(in_channels, 16, kernel_size=1)

    def _forward(self, x):
        branch1x1 = self.branch1x1(x)

        branch5x5 = self.branch5x5_1(x)
        branch5x5 = self.branch5x5_2(branch5x5)
        branch5x5 = self.branch5x5_3(branch5x5)

        branch5x5dbl = self.branch5x5dbl_1(x)
        branch5x5dbl = self.branch5x5dbl_2(branch5x5dbl)
        branch5x5dbl = self.branch5x5dbl_3(branch5x5dbl)
        branch5x5dbl = self.branch5x5dbl_4(branch5x5dbl)
        branch5x5dbl = self.branch5x5dbl_5(branch5x5dbl)

        branch_pool = nn.functional.avg_pool2d(x, kernel_size=3, stride=1, padding=1)
        branch_pool = self.branch_pool(branch_pool)

        outputs = [branch1x1, branch5x5, branch5x5dbl, branch_pool]

        return outputs

    def forward(self, x):
        outputs = self._forward(x)
        # print(outputs.shape)
        # torch.Size([2, 64, 6, 610])
        return torch.cat(outputs, 1)


if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = InceptionC(1, 16).to(device)
    print(summary(model, (1, 6, 610)))
